package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"strings"
)

var basicTypes = map[string]bool{
	"bool":       true,
	"string":     true,
	"int":        true,
	"int8":       true,
	"int16":      true,
	"int32":      true,
	"int64":      true,
	"uint":       true,
	"uint8":      true,
	"uint16":     true,
	"uint32":     true,
	"uint64":     true,
	"uintptr":    true,
	"byte":       true, // alias for uint8
	"rune":       true, // alias for int32
	"float32":    true,
	"float64":    true,
	"complex64":  true,
	"complex128": true,
}

var skipStructs []string

var (
	pbgoFile = flag.String("pbgo", "xxx.pb.go", "import pbgo file")
	pbgoHead = flag.String("head", "xxx", "package head")
	skip     = flag.String("skip", "C2L_|L2C_|C2S_|S2C_", "skip struct")
)

// 结构体字段类型信息
type FieldInfo struct {
	Name        string // 字段名称
	Type        string // 基本类型
	IsBasic     bool   // 是否基础类型
	IsMap       bool   // 是否map
	IsSlice     bool   // 是否slice
	IsStructPtr bool   // 是否*Struct类型
	KeyType     string // map key 类型
	ValueType   string // map value 或 slice 元素类型
}

// 获取字段类型描述
func parseFieldType(fieldName string, expr ast.Expr) (FieldInfo, error) {
	fi := FieldInfo{}
	switch t := expr.(type) {
	case *ast.Ident:
		fi.Type = t.Name
		fi.IsBasic = true
	case *ast.StarExpr:
		if ident, ok := t.X.(*ast.Ident); ok {
			fi.Type = ident.Name
			fi.IsStructPtr = true
		} else if sel, ok := t.X.(*ast.SelectorExpr); ok {
			// 支持导入的类型如 *pkg.Struct
			fi.Type = fmt.Sprintf("%s.%s", sel.X.(*ast.Ident).Name, sel.Sel.Name)
			fi.IsStructPtr = true
		}
	case *ast.MapType:
		key, err := parseFieldType(fieldName, t.Key)
		if err != nil {
			return fi, err
		}
		val, err := parseFieldType(fieldName, t.Value)
		if err != nil {
			return fi, err
		}
		fi.IsMap = true
		fi.KeyType = key.Type
		fi.ValueType = val.Type
	case *ast.ArrayType:
		elem, err := parseFieldType(fieldName, t.Elt)
		if err != nil {
			return fi, err
		}
		fi.IsSlice = true
		fi.ValueType = elem.Type
		// 处理 [][]byte 类型
		if fi.ValueType == "" {
			fi.Type = "[][]byte"
		}
	default:
		return fi, fmt.Errorf("unsupported field type: %T", expr)
	}
	return fi, nil
}

// 解析 Go 文件中定义的所有结构体
func parseStructsFromFile(filePath string) ([]*ast.StructType, []string, error) {
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, filePath, nil, parser.ParseComments)
	if err != nil {
		return nil, nil, err
	}

	var structs []*ast.StructType
	var names []string

	for _, decl := range node.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok || genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			skipFlag := false
			for _, st := range skipStructs {
				if strings.HasPrefix(typeSpec.Name.Name, st) {
					skipFlag = true
					break
				}
			}
			if skipFlag {
				continue
			}

			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			structs = append(structs, structType)
			names = append(names, typeSpec.Name.Name)
		}
	}
	return structs, names, nil
}

// 生成 Clone 函数
func generateCloneFunction(structName string, fields []*ast.Field) string {
	var code strings.Builder
	code.WriteString(fmt.Sprintf("func (m *%s) Clone() *%s {\n", structName, structName))
	code.WriteString("	if m != nil {\n")
	code.WriteString(fmt.Sprintf("		o := &%s{}\n", structName))

	for _, field := range fields {
		if len(field.Names) == 0 {
			continue
		}
		fieldName := field.Names[0].Name
		ft, _ := parseFieldType(fieldName, field.Type)

		if strings.HasPrefix(fieldName, "XXX_") {
			continue
		}

		switch {
		case ft.IsBasic:
			code.WriteString(fmt.Sprintf("		o.%s = m.%s\n", fieldName, fieldName))
		case ft.IsMap:
			code.WriteString(fmt.Sprintf("		if len(m.%s) > 0 {\n", fieldName))
			if _, exist := basicTypes[ft.ValueType]; exist {
				code.WriteString(fmt.Sprintf("			o.%s = make(map[%s]%s, len(m.%s))\n", fieldName, ft.KeyType, ft.ValueType, fieldName))
			} else {
				code.WriteString(fmt.Sprintf("			o.%s = make(map[%s]*%s, len(m.%s))\n", fieldName, ft.KeyType, ft.ValueType, fieldName))
			}

			code.WriteString(fmt.Sprintf("			for k, e := range m.%s {\n", fieldName))

			if _, exist := basicTypes[ft.ValueType]; !exist {
				code.WriteString(fmt.Sprintf("				o.%s[k] = e.Clone()\n", fieldName))
			} else {
				code.WriteString(fmt.Sprintf("				o.%s[k] = e\n", fieldName))
			}
			code.WriteString("			}\n")
			code.WriteString("		}\n")
		case ft.IsSlice:
			code.WriteString(fmt.Sprintf("		if len(m.%s) > 0 {\n", fieldName))
			if ft.Type == "[][]byte" {
				code.WriteString(fmt.Sprintf("			o.%s = make([][]byte, len(m.%s))\n", fieldName, fieldName))
				code.WriteString(fmt.Sprintf("			copy(o.%s, m.%s)\n", fieldName, fieldName))
			} else if _, exist := basicTypes[ft.ValueType]; exist {
				code.WriteString(fmt.Sprintf("			o.%s = make([]%s, len(m.%s))\n", fieldName, ft.ValueType, fieldName))
				code.WriteString(fmt.Sprintf("			copy(o.%s, m.%s)\n", fieldName, fieldName))
			} else {
				code.WriteString(fmt.Sprintf("			o.%s = make([]*%s, len(m.%s))\n", fieldName, ft.ValueType, fieldName))
				code.WriteString(fmt.Sprintf("			for k, e := range m.%s {\n", fieldName))
				code.WriteString(fmt.Sprintf("				o.%s[k] = e.Clone()\n", fieldName))
				code.WriteString("			}\n")
			}
			code.WriteString("		}\n")
		case ft.IsStructPtr:
			code.WriteString(fmt.Sprintf("		o.%s = m.%s.Clone()\n", fieldName, fieldName))
		}
	}

	code.WriteString("		return o\n")
	code.WriteString("	}\n")
	code.WriteString("	return nil\n")
	code.WriteString("}\n\n")
	return code.String()
}

// 生成 FixNilFields 函数
func generateFixNilFieldsFunction(structName string, fields []*ast.Field) string {
	var code strings.Builder
	code.WriteString(fmt.Sprintf("func (m *%s) FixNilFields() {\n", structName))
	code.WriteString("	if m != nil {\n\n")

	for _, field := range fields {
		if len(field.Names) == 0 {
			continue
		}
		fieldName := field.Names[0].Name
		ft, _ := parseFieldType(fieldName, field.Type)

		if strings.HasPrefix(fieldName, "XXX_") {
			continue
		}

		switch {
		case ft.IsMap:
			code.WriteString(fmt.Sprintf("		if m.%s == nil {\n", fieldName))
			if _, exist := basicTypes[ft.ValueType]; exist {
				code.WriteString(fmt.Sprintf("			m.%s = make(map[%s]%s)\n", fieldName, ft.KeyType, ft.ValueType))
			} else {
				code.WriteString(fmt.Sprintf("			m.%s = make(map[%s]*%s)\n", fieldName, ft.KeyType, ft.ValueType))
			}
			code.WriteString("		}\n")
			if _, exist := basicTypes[ft.ValueType]; !exist {
				code.WriteString(fmt.Sprintf("		for _, itor := range m.%s {\n", fieldName))
				code.WriteString(fmt.Sprintf("			itor.FixNilFields()\n"))
				code.WriteString("		}\n")
			}
			code.WriteString("\n")
		case ft.IsSlice:
			if _, exist := basicTypes[ft.ValueType]; !exist && ft.Type != "[][]byte" {
				code.WriteString(fmt.Sprintf("		for _, itor := range m.%s {\n", fieldName))
				code.WriteString(fmt.Sprintf("			itor.FixNilFields()\n"))
				code.WriteString("		}\n\n")
			}
		case ft.IsStructPtr:
			code.WriteString(fmt.Sprintf("		if m.%s == nil {\n", fieldName))
			code.WriteString(fmt.Sprintf("			m.%s = &%s{}\n", fieldName, ft.Type))
			code.WriteString(fmt.Sprintf("		}\n"))
			code.WriteString(fmt.Sprintf("		m.%s.FixNilFields()\n\n", fieldName))
		}
	}

	code.WriteString("	}\n")
	code.WriteString("}\n")
	return code.String()
}

// 生成 Init 函数
func generateInitFunction(structName string) string {
	return fmt.Sprintf("func Init%s() *%s {\n	m := &%s{}\n	m.FixNilFields()\n	return m\n}\n",
		structName, structName, structName)
}

// 主函数
func main() {
	//inputPath := "..\\cl\\cl.pb.go"        // 输入文件路径
	//outputPath := "..\\cl\\cl.pb.clone.go" // 输出文件路径
	flag.Parse()
	skipStructs = strings.Split(*skip, "|")

	if !strings.HasSuffix(*pbgoFile, ".pb.go") {
		fmt.Println("输入文件不是 .pb.go 文件")
		return
	}

	outputPath := strings.TrimSuffix(*pbgoFile, "go") + "clone.go"
	// 解析输入文件中的结构体
	structs, names, err := parseStructsFromFile(*pbgoFile)
	if err != nil {
		fmt.Printf("解析文件失败: %v\n", err)
		return
	}

	// 生成代码
	var outputCode strings.Builder
	outputCode.WriteString(fmt.Sprintf("package %s\n\n", *pbgoHead)) //设置包名

	for i, s := range structs {
		name := names[i]
		outputCode.WriteString(generateCloneFunction(name, s.Fields.List))
		outputCode.WriteString(generateInitFunction(name))
		outputCode.WriteString(generateFixNilFieldsFunction(name, s.Fields.List))
	}

	// 写入输出文件
	err = os.WriteFile(outputPath, []byte(outputCode.String()), 0644)
	if err != nil {
		fmt.Printf("写clone.go文件失败: %v\n", err)
		return
	}
	fmt.Printf("已成功生成clone文件: %s\n", outputPath)
}
